import torch
import torchvision
import torchvision.transforms as transforms
import numpy as np
import matplotlib.pyplot as plt
import os

# 设置保存数据的目录
data_dir = './data'

# 1️⃣ 定义图像预处理
transform = transforms.ToTensor()

# 2️⃣ 下载 CIFAR-10 数据集到指定文件夹（./data）
trainset = torchvision.datasets.CIFAR10(root=data_dir, train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root=data_dir, train=False, download=True, transform=transform)

# 3️⃣ 将图像和标签转换为 NumPy 格式
def dataset_to_numpy(dataset):
    X = []
    y = []
    for img, label in dataset:
        X.append(img.numpy().transpose(1, 2, 0))  # CHW -> HWC
        y.append(label)
    return np.array(X), np.array(y)

X_train, y_train = dataset_to_numpy(trainset)
X_test, y_test = dataset_to_numpy(testset)

print(f"✅ 数据已下载到: {os.path.abspath(data_dir)}")
print(f"训练集形状: {X_train.shape}, 测试集形状: {X_test.shape}")

# 4️⃣ 显示前 8 张图像和类别标签
class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer',
               'dog', 'frog', 'horse', 'ship', 'truck']

plt.figure(figsize=(10, 4))
for i in range(8):
    plt.subplot(2, 4, i + 1)
    plt.imshow(X_train[i])
    plt.title(class_names[y_train[i]])
    plt.axis('off')
plt.tight_layout()
plt.show()
